import logging
import os
import folder_paths
import numpy as np
import torch
import gc  # 添加gc模块导入
from PIL import Image, ImageOps, ImageSequence

# 配置日志
logger = logging.getLogger('PSDModifier')

# 确保IMAGE_PATH_LIST类型已定义
if not hasattr(folder_paths, 'IMAGE_PATH_LIST'):
    folder_paths.IMAGE_PATH_LIST = "IMAGE_PATH_LIST"

class LoadImagesFromDirectory:
    """
    从文件路径列表加载图像节点
    接收图像文件路径列表，根据指定索引加载对应的图像文件，
    并输出图像数据、蒙版及文件信息
    """
    # 定义节点的输入类型
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "image_paths": (folder_paths.IMAGE_PATH_LIST, {}),  # 从ListImagesInDirectory节点接收的图像路径列表
                "image_index": ("INT", {"default": 0, "min": 0, "max": 9999, "step": 1}),  # 要加载的图像索引
            },
            "optional": {
                "loop_index": ("BOOLEAN", {"default": True}),  # 是否循环索引（当超过总数时回到开始）
                "force_clean_cache": ("BOOLEAN", {"default": False}),  # 强制清理缓存参数
            }
        }

    # 定义节点的输出类型
    RETURN_TYPES = ("IMAGE", "MASK", "STRING", "STRING", "STRING", "STRING", "INT", "INT")
    RETURN_NAMES = ("IMAGE", "MASK", "full_path", "filename", "extension", "filename_no_ext", "image_count", "current_index")
    FUNCTION = "load_images"
    CATEGORY = "image/utils"
    OUTPUT_IS_LIST = (False, False, False, False, False, False, False, False)  # 所有输出都是单个值

    def load_images(self, image_paths=None, image_index=0, loop_index=True, force_clean_cache=False):
        """
        从传入的图像路径列表中加载指定索引的图像，支持内存清理功能
        
        参数:
            image_paths: 图像文件路径列表
            image_index: 要加载的图像索引
            loop_index: 是否循环索引（当超过总数时回到开始）
            force_clean_cache: 是否强制清理缓存
        
        返回:
            tuple: 包含图像数据、蒙版、完整路径、文件名、扩展名、无扩展名文件名、图像总数和当前索引的元组
        """
        # 如果强制清理缓存，执行显式内存清理
        if force_clean_cache:
            self._clean_memory()
        
        # 准备返回值
        output_image = torch.zeros((1, 64, 64, 3), dtype=torch.float32)  # 默认空图像
        output_mask = torch.zeros((1, 64, 64), dtype=torch.float32)  # 默认空蒙版
        output_full_path = ""
        output_filename = ""
        output_extension = ""
        output_filename_no_ext = ""
        current_index = 0
        image_count = 0
        
        # 检查输入的路径列表
        if not image_paths or not isinstance(image_paths, list) or len(image_paths) == 0:
            logger.error("LoadImagesFromDirectory: 无效的图像路径列表")
            return (output_image, output_mask, output_full_path, output_filename, output_extension, output_filename_no_ext, image_count, current_index)
        
        # 处理索引
        image_count = len(image_paths)
        if image_count > 0:
            if loop_index:
                # 循环索引：当索引超过总数时回到开始
                current_index = image_index % image_count
            else:
                # 限制索引在有效范围内
                current_index = min(max(image_index, 0), image_count - 1)
            
            try:
                # 获取当前要加载的图像文件路径
                file_path = image_paths[current_index]
                
                # 验证文件是否存在
                if not os.path.exists(file_path) or not os.path.isfile(file_path):
                    logger.error(f"LoadImagesFromDirectory: 文件不存在或不是有效文件: {file_path}")
                    return (output_image, output_mask, output_full_path, output_filename, output_extension, output_filename_no_ext, image_count, current_index)
                
                # 加载图像
                img = Image.open(file_path)
                
                # 处理EXIF方向
                try:
                    img = ImageOps.exif_transpose(img)
                except Exception as e:
                    logger.warning(f"LoadImagesFromDirectory: 处理EXIF方向失败 {file_path}: {str(e)}")
                    
                # 转换为RGB格式
                if img.mode == 'RGBA':
                    # 提取Alpha通道作为蒙版
                    mask = np.array(img.getchannel('A')).astype(np.float32) / 255.0
                    mask = np.expand_dims(mask, axis=0)
                    # 转换为RGB
                    img = img.convert('RGB')
                else:
                    # 如果没有Alpha通道，创建黑色蒙版
                    img = img.convert('RGB')
                    mask = np.zeros((img.height, img.width), dtype=np.float32)
                    mask = np.expand_dims(mask, axis=0)
                    
                # 转换为numpy数组并规范化
                img_array = np.array(img).astype(np.float32) / 255.0
                
                # 将numpy数组转换为PyTorch张量并添加批次维度
                # 格式变为(1, H, W, 3)，与ComfyUI兼容
                output_image = torch.from_numpy(img_array).unsqueeze(0)
                output_mask = torch.from_numpy(mask)
                
                # 提取文件信息
                output_full_path = os.path.abspath(file_path)
                output_filename = os.path.basename(output_full_path)
                _, output_extension = os.path.splitext(output_filename)
                output_extension = output_extension.lower().lstrip('.')
                output_filename_no_ext = os.path.splitext(output_filename)[0]
                
                logger.info(f"LoadImagesFromDirectory: 成功加载图像: {output_full_path} (索引: {current_index})")
                
            except Exception as e:
                logger.error(f"LoadImagesFromDirectory: 加载图像失败 {file_path}: {str(e)}")
                # 在发生错误时也尝试清理内存
                self._clean_memory()
        
        logger.info(f"LoadImagesFromDirectory: 总共找到 {image_count} 张图像，当前输出索引: {current_index}")
        
        return (output_image, output_mask, output_full_path, output_filename, output_extension, output_filename_no_ext, image_count, current_index)
    
    def _clean_memory(self):
        """
        执行显式内存清理操作
        """
        try:
            # 清理Torch缓存
            if torch.cuda.is_available():
                torch.cuda.empty_cache()
                torch.cuda.ipc_collect()
            # 执行Python垃圾回收
            gc.collect()
            logger.info("LoadImagesFromDirectory: 已执行显式内存清理")
        except Exception as e:
            logger.error(f"LoadImagesFromDirectory: 内存清理过程中发生错误: {str(e)}")

class CombineStringData:
    """
    字符串数据组合节点
    接收原始字符串数据和额外字符串数据，将它们组合后输出
    支持自定义分隔符
    """
    # 定义节点的输入类型
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "original_data": ("STRING", {"default": "", "multiline": True}),  # 原始字符串数据
                "additional_data": ("STRING", {"default": "", "multiline": True}),  # 要添加的额外字符串数据
            },
            "optional": {
                "separator": ("STRING", {"default": " ", "multiline": False}),  # 组合数据时使用的分隔符
            }
        }

    # 定义节点的输出类型
    RETURN_TYPES = ("STRING",)
    RETURN_NAMES = ("combined_result",)
    FUNCTION = "combine_data"
    CATEGORY = "text/utils"
    OUTPUT_IS_LIST = (False,)

    def combine_data(self, original_data="", additional_data="", separator=" "):
        """
        组合原始数据和额外数据
        
        参数:
            original_data: 原始字符串数据
            additional_data: 要添加的额外字符串数据
            separator: 组合数据时使用的分隔符
        
        返回:
            tuple: 包含组合后字符串的元组
        """
        try:
            # 确保输入是字符串类型
            original_str = str(original_data)
            additional_str = str(additional_data)
            separator_str = str(separator)
            
            # 组合字符串
            # 如果任一数据为空，则直接返回非空数据
            if not original_str:
                combined_str = additional_str
            elif not additional_str:
                combined_str = original_str
            else:
                # 使用分隔符组合两个非空字符串
                combined_str = f"{original_str}{separator_str}{additional_str}"
            
            logger.info(f"CombineStringData: 成功组合字符串数据")
            return (combined_str,)
        except Exception as e:
            logger.error(f"CombineStringData: 组合字符串数据时发生错误: {str(e)}")
            # 在出错时返回原始数据，保持稳定性
            return (str(original_data),)

# 注册节点
NODE_CLASS_MAPPINGS = {
    "LoadImagesFromDirectory": LoadImagesFromDirectory,
    "CombineStringData": CombineStringData
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "LoadImagesFromDirectory": "Load Images From Directory",
    "CombineStringData": "Combine String Data"
}